{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "750a91b5",
   "metadata": {},
   "source": [
    "# STORM 개념을 도입한 연구를 위한 멀티 에이전트\n",
    "\n",
    "이 튜토리얼의 목적은 LangGraph를 활용하여 연구 자동화 시스템을 구축하는 방법에 대해서 다룹니다. \n",
    "\n",
    "연구는 종종 분석가에게 위임되는 노동 집약적인 작업입니다. AI는 이러한 연구 과정을 지원할 수 있는 상당한 잠재력을 가지고 있습니다. 이 튜토리얼에서는 사용자 맞춤형 AI 기반 연구 및 보고서 생성 워크플로우를 구축하는 방법을 다룹니다.\n",
    "\n",
    "이번 튜토리얼에서는 경량의 다중 에이전트 시스템을 구축하여 연구 과정을 맞춤화하는 것을 목표로 합니다. 사용자는 연구 주제를 제공하고, 시스템은 각 하위 주제에 집중하는 AI 분석가 팀을 생성합니다. \n",
    "\n",
    "이 과정에서 `Human-in-the-loop`를 사용하여 연구가 시작되기 전에 하위 주제를 세분화합니다.\n",
    "\n",
    "[STORM 논문](https://arxiv.org/abs/2402.14207)에 따르면, **유사한 주제 조회**와 **다양한 관점의 대화 시뮬레이션**을 통해 참고 출처 사용 빈도와 정보 밀도를 증가시킬 수 있습니다. \n",
    "\n",
    "**주로 다루는 내용**\n",
    "- **LangGraph의 주요 테마**: Memory, Human-in-the-loop, Controllability\n",
    "- **연구 자동화의 목표**: 사용자 맞춤형 연구 프로세스 구축\n",
    "- **소스 선택**: 연구를 위한 입력 소스 선택\n",
    "- **계획**: 주제 제공 및 AI 분석가 팀 생성\n",
    "- **LLM 활용**: 전문가 AI와의 심층 인터뷰\n",
    "- **연구 과정**: 병렬로 정보 수집 및 인터뷰 수행\n",
    "- **출력 형식**: 최종 보고서로 통합된 통찰력\n",
    "- **설정**: 환경 설정 및 API 키 설정\n",
    "- **분석가 생성**: Human-In-The-Loop를 통한 분석가 생성 및 검토\n",
    "- **인터뷰 수행**: 질문 생성 및 답변 수집\n",
    "- **병렬 인터뷰**: Map-Reduce를 통한 인터뷰 병렬화\n",
    "- **최종 보고서 작성**: 보고서의 서론 및 결론 작성"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6339db37",
   "metadata": {},
   "source": [
    "이번 튜토리얼에서는 다음의 세 가지 테마를 다룹니다.\n",
    "\n",
    "- **Memory**\n",
    "- **Human-in-the-loop**\n",
    "- **Controllability**\n",
    "\n",
    "이제 이러한 개념을 결합하여 AI의 가장 인기 있는 응용 분야 중 하나인 연구 자동화를 다루겠습니다. \n",
    "\n",
    "연구는 종종 분석가에게 위임되는 노동 집약적인 작업입니다. AI는 이러한 연구 과정을 지원할 수 있는 상당한 잠재력을 가지고 있습니다. 그러나 연구는 맞춤화가 필요합니다. 원시 LLM 출력은 실제 의사 결정 워크플로우에 적합하지 않은 경우가 많습니다.\n",
    "\n",
    "맞춤형 AI 기반 [연구 및 보고서 생성](https://jxnl.co/writing/2024/06/05/predictions-for-the-future-of-rag/#reports-over-rag) 워크플로우는 이를 해결할 수 있는 유망한 방법입니다.\n",
    "\n",
    "![langgraph-storm-concept](./assets/langgraph-storm-concept.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9c5b8ea",
   "metadata": {},
   "source": [
    "## 환경 설정"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d41258fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# API 키를 환경변수로 관리하기 위한 설정 파일\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "# API 키 정보 로드\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2dad8507",
   "metadata": {},
   "outputs": [],
   "source": [
    "# LangSmith 추적을 설정합니다. https://smith.langchain.com\n",
    "# !pip install -qU langchain-teddynote\n",
    "from langchain_teddynote import logging\n",
    "\n",
    "# 프로젝트 이름을 입력합니다.\n",
    "logging.langsmith(\"CH17-LangGraph-Use-Cases\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "889da53e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_teddynote.models import get_model_name, LLMs\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "# 최신 모델 가져오기\n",
    "GPT4o = get_model_name(LLMs.GPT4o)\n",
    "\n",
    "# 모델 초기화\n",
    "llm = ChatOpenAI(model=GPT4o)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7017f222",
   "metadata": {},
   "source": [
    "## 분석가 생성: Human-In-The-Loop\n",
    "\n",
    "- **분석가 생성**: `Human-In-The-Loop`를 활용하여 분석가를 생성하고 검토합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "be7f039b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "from typing_extensions import TypedDict\n",
    "from pydantic import BaseModel, Field\n",
    "from langchain_teddynote.graphs import visualize_graph\n",
    "\n",
    "\n",
    "# 분석가의 속성과 메타데이터를 정의하는 클래스\n",
    "class Analyst(BaseModel):\n",
    "    # 주요 소속 정보\n",
    "    affiliation: str = Field(\n",
    "        description=\"Primary affiliation of the analyst.\",\n",
    "    )\n",
    "    # 이름\n",
    "    name: str = Field(description=\"Name of the analyst.\")\n",
    "\n",
    "    # 역할\n",
    "    role: str = Field(\n",
    "        description=\"Role of the analyst in the context of the topic.\",\n",
    "    )\n",
    "    # 중점, 우려 사항 및 동기에 대한 설명\n",
    "    description: str = Field(\n",
    "        description=\"Description of the analyst focus, concerns, and motives.\",\n",
    "    )\n",
    "\n",
    "    # 분석가의 인적 정보를 문자열로 반환하는 속성\n",
    "    @property\n",
    "    def persona(self) -> str:\n",
    "        return f\"Name: {self.name}\\nRole: {self.role}\\nAffiliation: {self.affiliation}\\nDescription: {self.description}\\n\"\n",
    "\n",
    "\n",
    "# 분석가들의 집합\n",
    "class Perspectives(BaseModel):\n",
    "    # 분석가 목록\n",
    "    analysts: List[Analyst] = Field(\n",
    "        description=\"Comprehensive list of analysts with their roles and affiliations.\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4845a360",
   "metadata": {},
   "source": [
    "다음은 Analyst 클래스를 통해 생성된 분석가들의 집합을 추적하는 상태를 정의합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8dba2722",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 상태 정의\n",
    "class GenerateAnalystsState(TypedDict):\n",
    "    # 연구 주제\n",
    "    topic: str\n",
    "    # 생성할 분석가의 최대 수\n",
    "    max_analysts: int\n",
    "    # 사람 피드백\n",
    "    human_analyst_feedback: str\n",
    "    # 분석가 목록\n",
    "    analysts: List[Analyst]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dba448ac",
   "metadata": {},
   "source": [
    "## 분석가(Analyst) 생성 노드 정의"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cd092f4",
   "metadata": {},
   "source": [
    "다음으로는 분석가(Analyst) 생성 노드를 정의하겠습니다.\n",
    "\n",
    "아래 코드는 주어진 연구 주제에 대해 다양한 분석가를 생성하는 로직을 구현합니다. 각 분석가는 고유한 역할과 소속을 가지며, 주제에 대한 전문적인 관점을 제공합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5cd7ae07",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph import END\n",
    "from langchain_core.messages import HumanMessage, SystemMessage\n",
    "\n",
    "# 분석가 생성 프롬프트\n",
    "analyst_instructions = \"\"\"You are tasked with creating a set of AI analyst personas. \n",
    "\n",
    "Follow these instructions carefully:\n",
    "1. First, review the research topic:\n",
    "\n",
    "{topic}\n",
    "        \n",
    "2. Examine any editorial feedback that has been optionally provided to guide creation of the analysts: \n",
    "        \n",
    "{human_analyst_feedback}\n",
    "    \n",
    "3. Determine the most interesting themes based upon documents and / or feedback above.\n",
    "                    \n",
    "4. Pick the top {max_analysts} themes.\n",
    "\n",
    "5. Assign one analyst to each theme.\"\"\"\n",
    "\n",
    "\n",
    "# 분석가 생성 노드\n",
    "def create_analysts(state: GenerateAnalystsState):\n",
    "    \"\"\"분석가 페르소나를 생성하는 함수\"\"\"\n",
    "\n",
    "    topic = state[\"topic\"]\n",
    "    max_analysts = state[\"max_analysts\"]\n",
    "    human_analyst_feedback = state.get(\"human_analyst_feedback\", \"\")\n",
    "\n",
    "    # LLM에 구조화된 출력 형식을 적용\n",
    "    structured_llm = llm.with_structured_output(Perspectives)\n",
    "\n",
    "    # 분석가 생성을 위한 시스템 프롬프트 구성\n",
    "    system_message = analyst_instructions.format(\n",
    "        topic=topic,\n",
    "        human_analyst_feedback=human_analyst_feedback,\n",
    "        max_analysts=max_analysts,\n",
    "    )\n",
    "\n",
    "    # LLM을 호출하여 분석가 페르소나 생성\n",
    "    analysts = structured_llm.invoke(\n",
    "        [SystemMessage(content=system_message)]\n",
    "        + [HumanMessage(content=\"Generate the set of analysts.\")]\n",
    "    )\n",
    "\n",
    "    # 생성된 분석가 목록을 상태에 저장\n",
    "    return {\"analysts\": analysts.analysts}\n",
    "\n",
    "\n",
    "# 사용자 피드백 노드(상태 업데이트를 진행할 예정이므로, 내용은 비워 두어도 무방)\n",
    "def human_feedback(state: GenerateAnalystsState):\n",
    "    \"\"\"사용자 피드백을 받기 위한 중단점 노드\"\"\"\n",
    "    pass\n",
    "\n",
    "\n",
    "# 인간 피드백 여부에 따라 워크플로우의 다음 단계를 결정하는 함수\n",
    "def should_continue(state: GenerateAnalystsState):\n",
    "    \"\"\"워크플로우의 다음 단계를 결정하는 함수\"\"\"\n",
    "\n",
    "    human_analyst_feedback = state.get(\"human_analyst_feedback\", None)\n",
    "    if human_analyst_feedback:\n",
    "        return \"create_analysts\"\n",
    "\n",
    "    return END"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22196430",
   "metadata": {},
   "source": [
    "## 그래프 생성\n",
    "\n",
    "이제 분석가 생성 그래프를 생성합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f33720ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph import START, END, StateGraph\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "from langchain_core.messages import SystemMessage\n",
    "\n",
    "# 그래프 생성\n",
    "builder = StateGraph(GenerateAnalystsState)\n",
    "\n",
    "# 노드 추가\n",
    "builder.add_node(\"create_analysts\", create_analysts)\n",
    "builder.add_node(\"human_feedback\", human_feedback)\n",
    "\n",
    "# 엣지 연결\n",
    "builder.add_edge(START, \"create_analysts\")\n",
    "builder.add_edge(\"create_analysts\", \"human_feedback\")\n",
    "\n",
    "# 조건부 엣지 추가: 사람 피드백이 있을 경우 다시 분석가 생성 노드로 돌아갑니다.\n",
    "builder.add_conditional_edges(\n",
    "    \"human_feedback\", should_continue, [\"create_analysts\", END]\n",
    ")\n",
    "\n",
    "# 메모리 생성\n",
    "memory = MemorySaver()\n",
    "\n",
    "# 그래프 컴파일(중단점 설정)\n",
    "graph = builder.compile(interrupt_before=[\"human_feedback\"], checkpointer=memory)\n",
    "\n",
    "# 그래프 시각화\n",
    "visualize_graph(graph)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9b96538",
   "metadata": {},
   "source": [
    "## 분석가 생성을 위한 그래프 실행"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06ddccc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.runnables import RunnableConfig\n",
    "from langchain_teddynote.messages import random_uuid, invoke_graph\n",
    "\n",
    "config = RunnableConfig(\n",
    "    recursion_limit=10,\n",
    "    configurable={\"thread_id\": random_uuid()},\n",
    ")\n",
    "\n",
    "# 분석가 수 설정\n",
    "max_analysts = 3\n",
    "\n",
    "# 연구 주제 설정\n",
    "topic = \"Modular RAG 가 기존의 Naive RAG 와 어떤 차이가 있는지와 production level 에서 사용하는 이점\"\n",
    "\n",
    "# 입력 데이터 설정\n",
    "inputs = {\n",
    "    \"topic\": topic,\n",
    "    \"max_analysts\": max_analysts,\n",
    "}\n",
    "\n",
    "# 그래프 실행\n",
    "invoke_graph(graph, inputs, config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "623117de",
   "metadata": {},
   "source": [
    "`__interrupt__` 가 출력되면 인간의 피드백을 받을 준비가 된 것입니다.\n",
    "\n",
    "이제 아래의 상태를 가져와서 인간 피드백을 제공합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8fe74eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 그래프의 현재 상태 가져오기\n",
    "state = graph.get_state(config)\n",
    "\n",
    "# 다음 실행할 노드 확인\n",
    "print(state.next)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56d6ca4e",
   "metadata": {},
   "source": [
    "`update_state()` 를 통해 인간 피드백을 주입합니다. 이때 `human_analyst_feedback` 키에 피드백 내용을 저장합니다.\n",
    "\n",
    "또한 `as_node` 인자를 통해 피드백을 받을 노드를 지정합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efba1d3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 그래프 상태를 업데이트하여 human_feedback 노드의 역할 수행\n",
    "graph.update_state(\n",
    "    config,\n",
    "    {\n",
    "        \"human_analyst_feedback\": \"Add in someone named Teddy Lee from a startup to add an entrepreneur perspective\"\n",
    "    },\n",
    "    as_node=\"human_feedback\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52b22e9d",
   "metadata": {},
   "source": [
    "`None` 값을 입력으로 주게 되면, 이어서 그래프가 진행됩니다.\n",
    "\n",
    "**참고**\n",
    "\n",
    "- 재개하고자 할 때는 입력에 `None` 값을 할당하여 그래프를 재개합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65011c31",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 이어서 진행\n",
    "invoke_graph(graph, None, config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf393cce",
   "metadata": {},
   "source": [
    "다시 `__interrupt__` 가 출력되면 인간의 피드백을 받을 준비가 된 것입니다.\n",
    "\n",
    "이전의 방식과 동일하게 다시 인간 피드백을 제공하여 생성된 분석가의 페르소나를 조정하는 것도 가능합니다.\n",
    "\n",
    "하지만, 추가 피드백이 없을 경우 `None` 값을 할당하여 분석가 생성 작업을 종료할 수 있습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab512d7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 추가 피드백이 없을 경우 None 값을 할당하여 상태 업데이트\n",
    "human_feedback_input = None\n",
    "\n",
    "# 그래프 상태를 업데이트하여 human_feedback 노드의 역할 수행\n",
    "graph.update_state(\n",
    "    config, {\"human_analyst_feedback\": human_feedback_input}, as_node=\"human_feedback\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "35618b90",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 이어서 진행\n",
    "invoke_graph(graph, None, config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "710300b6",
   "metadata": {},
   "source": [
    "최종 결과를 출력합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d759708e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 그래프의 최종 상태 가져오기\n",
    "final_state = graph.get_state(config)\n",
    "\n",
    "# 최종 상태에서 생성된 분석가 목록 가져오기\n",
    "analysts = final_state.values.get(\"analysts\")\n",
    "\n",
    "# 생성된 분석가 수 출력\n",
    "print(f\"생성된 분석가 수: {len(analysts)}\", end=\"\\n================================\\n\")\n",
    "\n",
    "# 각 분석가의 페르소나 출력\n",
    "for analyst in analysts:\n",
    "    print(analyst.persona)\n",
    "    print(\"- \" * 30)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c61b611",
   "metadata": {},
   "source": [
    "`final_state.next` 는 그래프의 다음 실행할 노드를 나타냅니다. 여기서는 모든 작업이 마무리 되었기 때문에 빈 `tuple` 이 출력됩니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f92144e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 그래프의 다음 실행할 노드 상태 가져오기\n",
    "print(final_state.next)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6e14e5c",
   "metadata": {},
   "source": [
    "## 인터뷰 수행\n",
    "\n",
    "### 질문 생성\n",
    "\n",
    "- 분석가는 전문가에게 질문을 제시합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "089de179",
   "metadata": {},
   "outputs": [],
   "source": [
    "import operator\n",
    "from typing import Annotated\n",
    "from langgraph.graph import MessagesState\n",
    "\n",
    "\n",
    "# 인터뷰 상태 정의\n",
    "class InterviewState(MessagesState):\n",
    "    # 대화 턴수\n",
    "    max_num_turns: int\n",
    "    # 소스 문서를 포함하는 컨텍스트 리스트\n",
    "    context: Annotated[list, operator.add]\n",
    "    # 지정된 분석가\n",
    "    analyst: Analyst\n",
    "    # 인터뷰 내용을 저장하는 문자열\n",
    "    interview: str\n",
    "    # 보고서 섹션 리스트\n",
    "    sections: list\n",
    "\n",
    "\n",
    "# 검색 쿼리 데이터 클래스 정의\n",
    "class SearchQuery(BaseModel):\n",
    "    search_query: str = Field(None, description=\"Search query for retrieval.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1bd2a54",
   "metadata": {},
   "source": [
    "다음으로는 인터뷰 질문을 생성하는 노드를 정의합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "a74a90ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "question_instructions = \"\"\"You are an analyst tasked with interviewing an expert to learn about a specific topic. \n",
    "\n",
    "Your goal is boil down to interesting and specific insights related to your topic.\n",
    "\n",
    "1. Interesting: Insights that people will find surprising or non-obvious.\n",
    "        \n",
    "2. Specific: Insights that avoid generalities and include specific examples from the expert.\n",
    "\n",
    "Here is your topic of focus and set of goals: {goals}\n",
    "        \n",
    "Begin by introducing yourself using a name that fits your persona, and then ask your question.\n",
    "\n",
    "Continue to ask questions to drill down and refine your understanding of the topic.\n",
    "        \n",
    "When you are satisfied with your understanding, complete the interview with: \"Thank you so much for your help!\"\n",
    "\n",
    "Remember to stay in character throughout your response, reflecting the persona and goals provided to you.\"\"\"\n",
    "\n",
    "\n",
    "# 질문을 생성하는 노드 정의\n",
    "def generate_question(state: InterviewState):\n",
    "    # 상태에서 분석가와 메시지 가져오기\n",
    "    analyst = state[\"analyst\"]\n",
    "    messages = state[\"messages\"]\n",
    "\n",
    "    # 질문 생성\n",
    "    # 분석가의 목표를 기반으로 시스템 메시지 생성\n",
    "    system_message = question_instructions.format(goals=analyst.persona)\n",
    "    # LLM을 사용하여 질문 생성\n",
    "    question = llm.invoke([SystemMessage(content=system_message)] + messages)\n",
    "\n",
    "    # 상태에 메시지 기록\n",
    "    return {\"messages\": [question]}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "064303de",
   "metadata": {},
   "source": [
    "## 도구 정의\n",
    "\n",
    "전문가는 여러 소스로부터 정보를 병렬로 수집하여 질문에 답변합니다.\n",
    "\n",
    "웹 문서 스크래핑, VectorDB, 웹 검색, 위키피디아 검색 등 다양한 도구를 사용할 수 있습니다.\n",
    "\n",
    "이 튜토리얼에서는 Arxiv, Tavily 검색을 사용합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "592a8d12",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 웹 검색 도구 초기화\n",
    "from langchain_teddynote.tools.tavily import TavilySearch\n",
    "\n",
    "# 웹 검색을 위한 TavilySearch 인스턴스 생성\n",
    "tavily_search = TavilySearch(max_results=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd85cc55",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.retrievers import ArxivRetriever\n",
    "\n",
    "# Arxiv 검색을 위한 ArxivRetriever 인스턴스 생성\n",
    "arxiv_retriever = ArxivRetriever(\n",
    "    load_max_docs=3,\n",
    "    load_all_available_meta=True,\n",
    "    get_full_documents=True,\n",
    ")\n",
    "\n",
    "# 검색 결과 출력\n",
    "arxiv_search_results = arxiv_retriever.invoke(\"Modular RAG vs Naive RAG\")\n",
    "print(arxiv_search_results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03485bb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Arxiv 메타데이터 출력\n",
    "arxiv_search_results[0].metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fc31182",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Arxiv 내용 출력\n",
    "print(arxiv_search_results[0].page_content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad6294ff",
   "metadata": {},
   "source": [
    "문서 검색결과를 포맷팅 출력합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "d0475512",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 문서 검색 결과를 포맷팅\n",
    "formatted_search_docs = \"\\n\\n---\\n\\n\".join(\n",
    "    [\n",
    "        f'<Document source=\"{doc.metadata[\"entry_id\"]}\" date=\"{doc.metadata.get(\"Published\", \"\")}\" authors=\"{doc.metadata.get(\"Authors\", \"\")}\"/>\\n<Title>\\n{doc.metadata[\"Title\"]}\\n</Title>\\n\\n<Summary>\\n{doc.metadata[\"Summary\"]}\\n</Summary>\\n\\n<Content>\\n{doc.page_content}\\n</Content>\\n</Document>'\n",
    "        for doc in arxiv_search_results\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abc66263",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(formatted_search_docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5439f615",
   "metadata": {},
   "source": [
    "## 노드 생성\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "d26cd924",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import get_buffer_string\n",
    "from langchain_core.messages import SystemMessage, HumanMessage, AIMessage\n",
    "\n",
    "# 검색 쿼리 작성\n",
    "search_instructions = SystemMessage(\n",
    "    content=f\"\"\"You will be given a conversation between an analyst and an expert. \n",
    "\n",
    "Your goal is to generate a well-structured query for use in retrieval and / or web-search related to the conversation.\n",
    "        \n",
    "First, analyze the full conversation.\n",
    "\n",
    "Pay particular attention to the final question posed by the analyst.\n",
    "\n",
    "Convert this final question into a well-structured web search query\"\"\"\n",
    ")\n",
    "\n",
    "\n",
    "# 웹 검색 수행 함수 정의\n",
    "def search_web(state: InterviewState):\n",
    "    \"\"\"웹 검색을 통한 문서 검색\"\"\"\n",
    "\n",
    "    # 검색 쿼리 생성\n",
    "    structured_llm = llm.with_structured_output(SearchQuery)\n",
    "    search_query = structured_llm.invoke([search_instructions] + state[\"messages\"])\n",
    "\n",
    "    # 검색 수행\n",
    "    search_docs = tavily_search.invoke(search_query.search_query)\n",
    "\n",
    "    # 검색 결과 형식 지정\n",
    "    formatted_search_docs = \"\\n\\n---\\n\\n\".join(\n",
    "        [\n",
    "            f'<Document href=\"{doc[\"url\"]}\"/>\\n{doc[\"content\"]}\\n</Document>'\n",
    "            for doc in search_docs\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    return {\"context\": [formatted_search_docs]}\n",
    "\n",
    "\n",
    "# Arxiv 검색 노드 생성\n",
    "def search_arxiv(state: InterviewState):\n",
    "    \"\"\"Arxiv 검색 노드\"\"\"\n",
    "\n",
    "    # 검색 쿼리 생성\n",
    "    structured_llm = llm.with_structured_output(SearchQuery)\n",
    "    search_query = structured_llm.invoke([search_instructions] + state[\"messages\"])\n",
    "\n",
    "    try:\n",
    "        # 검색 수행\n",
    "        arxiv_search_results = arxiv_retriever.invoke(\n",
    "            search_query.search_query,\n",
    "            load_max_docs=2,\n",
    "            load_all_available_meta=True,\n",
    "            get_full_documents=True,\n",
    "        )\n",
    "\n",
    "        # 검색 결과 형식 지정\n",
    "        formatted_search_docs = \"\\n\\n---\\n\\n\".join(\n",
    "            [\n",
    "                f'<Document source=\"{doc.metadata[\"entry_id\"]}\" date=\"{doc.metadata.get(\"Published\", \"\")}\" authors=\"{doc.metadata.get(\"Authors\", \"\")}\"/>\\n<Title>\\n{doc.metadata[\"Title\"]}\\n</Title>\\n\\n<Summary>\\n{doc.metadata[\"Summary\"]}\\n</Summary>\\n\\n<Content>\\n{doc.page_content}\\n</Content>\\n</Document>'\n",
    "                for doc in arxiv_search_results\n",
    "            ]\n",
    "        )\n",
    "\n",
    "        return {\"context\": [formatted_search_docs]}\n",
    "    except Exception as e:\n",
    "        print(f\"Arxiv 검색 중 오류 발생: {str(e)}\")\n",
    "        return {\n",
    "            \"context\": [\"<Error>Arxiv 검색 결과를 가져오는데 실패했습니다.</Error>\"]\n",
    "        }\n",
    "\n",
    "\n",
    "answer_instructions = \"\"\"You are an expert being interviewed by an analyst.\n",
    "\n",
    "Here is analyst area of focus: {goals}. \n",
    "        \n",
    "You goal is to answer a question posed by the interviewer.\n",
    "\n",
    "To answer question, use this context:\n",
    "        \n",
    "{context}\n",
    "\n",
    "When answering questions, follow these guidelines:\n",
    "        \n",
    "1. Use only the information provided in the context. \n",
    "        \n",
    "2. Do not introduce external information or make assumptions beyond what is explicitly stated in the context.\n",
    "\n",
    "3. The context contain sources at the topic of each individual document.\n",
    "\n",
    "4. Include these sources your answer next to any relevant statements. For example, for source # 1 use [1]. \n",
    "\n",
    "5. List your sources in order at the bottom of your answer. [1] Source 1, [2] Source 2, etc\n",
    "        \n",
    "6. If the source is: <Document source=\"assistant/docs/llama3_1.pdf\" page=\"7\"/>' then just list: \n",
    "        \n",
    "[1] assistant/docs/llama3_1.pdf, page 7 \n",
    "        \n",
    "And skip the addition of the brackets as well as the Document source preamble in your citation.\"\"\"\n",
    "\n",
    "\n",
    "# 질문에 대한 답변 생성 함수 정의\n",
    "def generate_answer(state: InterviewState):\n",
    "    \"\"\"질문에 대한 답변 생성 노드\"\"\"\n",
    "\n",
    "    # 상태에서 분석가와 메시지 가져오기\n",
    "    analyst = state[\"analyst\"]\n",
    "    messages = state[\"messages\"]\n",
    "    context = state[\"context\"]\n",
    "\n",
    "    # 질문에 대한 답변 생성\n",
    "    system_message = answer_instructions.format(goals=analyst.persona, context=context)\n",
    "    answer = llm.invoke([SystemMessage(content=system_message)] + messages)\n",
    "\n",
    "    # 메시지를 전문가의 답변으로 명명\n",
    "    answer.name = \"expert\"\n",
    "\n",
    "    # 상태에 메시지 추가\n",
    "    return {\"messages\": [answer]}\n",
    "\n",
    "\n",
    "# 인터뷰 저장 함수 정의\n",
    "def save_interview(state: InterviewState):\n",
    "    \"\"\"인터뷰 저장\"\"\"\n",
    "\n",
    "    # 메시지 가져오기\n",
    "    messages = state[\"messages\"]\n",
    "\n",
    "    # 인터뷰를 문자열로 변환\n",
    "    interview = get_buffer_string(messages)\n",
    "\n",
    "    # 인터뷰 키에 저장\n",
    "    return {\"interview\": interview}\n",
    "\n",
    "\n",
    "# 메시지 라우팅 함수 정의\n",
    "def route_messages(state: InterviewState, name: str = \"expert\"):\n",
    "    \"\"\"질문과 답변 사이의 라우팅\"\"\"\n",
    "\n",
    "    # 메시지 가져오기\n",
    "    messages = state[\"messages\"]\n",
    "    max_num_turns = state.get(\"max_num_turns\", 2)\n",
    "\n",
    "    # 전문가의 답변 수 확인\n",
    "    num_responses = len(\n",
    "        [m for m in messages if isinstance(m, AIMessage) and m.name == name]\n",
    "    )\n",
    "\n",
    "    # 전문가가 최대 턴 수 이상 답변한 경우 종료\n",
    "    if num_responses >= max_num_turns:\n",
    "        return \"save_interview\"\n",
    "\n",
    "    # 이 라우터는 각 질문-답변 쌍 후에 실행됨\n",
    "    # 논의 종료를 신호하는 마지막 질문 가져오기\n",
    "    last_question = messages[-2]\n",
    "\n",
    "    if \"Thank you so much for your help\" in last_question.content:\n",
    "        return \"save_interview\"\n",
    "    return \"ask_question\"\n",
    "\n",
    "\n",
    "# 세션 작성 지시사항\n",
    "section_writer_instructions = \"\"\"You are an expert technical writer. \n",
    "\n",
    "Your task is to create a detailed and comprehensive section of a report, thoroughly analyzing a set of source documents.\n",
    "This involves extracting key insights, elaborating on relevant points, and providing in-depth explanations to ensure clarity and understanding. Your writing should include necessary context, supporting evidence, and examples to enhance the reader's comprehension. Maintain a logical and well-organized structure, ensuring that all critical aspects are covered in detail and presented in a professional tone.\n",
    "\n",
    "Please follow these instructions:\n",
    "1. Analyze the content of the source documents: \n",
    "- The name of each source document is at the start of the document, with the <Document tag.\n",
    "        \n",
    "2. Create a report structure using markdown formatting:\n",
    "- Use ## for the section title\n",
    "- Use ### for sub-section headers\n",
    "        \n",
    "3. Write the report following this structure:\n",
    "a. Title (## header)\n",
    "b. Summary (### header)\n",
    "c. Comprehensive analysis (### header)\n",
    "d. Sources (### header)\n",
    "\n",
    "4. Make your title engaging based upon the focus area of the analyst: \n",
    "{focus}\n",
    "\n",
    "5. For the summary section:\n",
    "- Set up summary with general background / context related to the focus area of the analyst\n",
    "- Emphasize what is novel, interesting, or surprising about insights gathered from the interview\n",
    "- Create a numbered list of source documents, as you use them\n",
    "- Do not mention the names of interviewers or experts\n",
    "- Aim for approximately 400 words maximum\n",
    "- Use numbered sources in your report (e.g., [1], [2]) based on information from source documents\n",
    "\n",
    "6. For the Comprehensive analysis section:\n",
    "- Provide a detailed examination of the information from the source documents.\n",
    "- Break down complex ideas into digestible segments, ensuring a logical flow of ideas.\n",
    "- Use sub-sections where necessary to cover multiple perspectives or dimensions of the analysis.\n",
    "- Support your analysis with data, direct quotes, and examples from the source documents.\n",
    "- Clearly explain the relevance of each point to the overall focus of the report.\n",
    "- Use bullet points or numbered lists for clarity when presenting multiple related ideas.\n",
    "- Ensure the tone remains professional and objective, avoiding bias or unsupported opinions.\n",
    "- Aim for at least 800 words to ensure the analysis is thorough.\n",
    "\n",
    "7. In the Sources section:\n",
    "- Include all sources used in your report\n",
    "- Provide full links to relevant websites or specific document paths\n",
    "- Separate each source by a newline. Use two spaces at the end of each line to create a newline in Markdown.\n",
    "- It will look like:\n",
    "\n",
    "### Sources\n",
    "[1] Link or Document name\n",
    "[2] Link or Document name\n",
    "\n",
    "8. Be sure to combine sources. For example this is not correct:\n",
    "\n",
    "[3] https://ai.meta.com/blog/meta-llama-3-1/\n",
    "[4] https://ai.meta.com/blog/meta-llama-3-1/\n",
    "\n",
    "There should be no redundant sources. It should simply be:\n",
    "\n",
    "[3] https://ai.meta.com/blog/meta-llama-3-1/\n",
    "        \n",
    "9. Final review:\n",
    "- Ensure the report follows the required structure\n",
    "- Include no preamble before the title of the report\n",
    "- Check that all guidelines have been followed\"\"\"\n",
    "\n",
    "\n",
    "# 섹션 작성 함수 정의\n",
    "def write_section(state: InterviewState):\n",
    "    \"\"\"질문에 대한 답변 생성 노드\"\"\"\n",
    "\n",
    "    # 상태에서 컨텍스트, 분석가 가져오기\n",
    "    context = state[\"context\"]\n",
    "    analyst = state[\"analyst\"]\n",
    "\n",
    "    # 섹션 작성을 위한 시스템 프롬프트 정의\n",
    "    system_message = section_writer_instructions.format(focus=analyst.description)\n",
    "    section = llm.invoke(\n",
    "        [SystemMessage(content=system_message)]\n",
    "        + [HumanMessage(content=f\"Use this source to write your section: {context}\")]\n",
    "    )\n",
    "\n",
    "    # 상태에 섹션 추가\n",
    "    return {\"sections\": [section.content]}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7fb721d",
   "metadata": {},
   "source": [
    "## 인터뷰 그래프 생성\n",
    "\n",
    "인터뷰를 수행하는 그래프를 정의하고 실행합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f74e1e7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph import StateGraph\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "\n",
    "# 노드 및 엣지 추가\n",
    "interview_builder = StateGraph(InterviewState)\n",
    "interview_builder.add_node(\"ask_question\", generate_question)\n",
    "interview_builder.add_node(\"search_web\", search_web)\n",
    "interview_builder.add_node(\"search_arxiv\", search_arxiv)\n",
    "interview_builder.add_node(\"answer_question\", generate_answer)\n",
    "interview_builder.add_node(\"save_interview\", save_interview)\n",
    "interview_builder.add_node(\"write_section\", write_section)\n",
    "\n",
    "# 흐름 설정\n",
    "interview_builder.add_edge(START, \"ask_question\")\n",
    "interview_builder.add_edge(\"ask_question\", \"search_web\")\n",
    "interview_builder.add_edge(\"ask_question\", \"search_arxiv\")\n",
    "interview_builder.add_edge(\"search_web\", \"answer_question\")\n",
    "interview_builder.add_edge(\"search_arxiv\", \"answer_question\")\n",
    "interview_builder.add_conditional_edges(\n",
    "    \"answer_question\", route_messages, [\"ask_question\", \"save_interview\"]\n",
    ")\n",
    "interview_builder.add_edge(\"save_interview\", \"write_section\")\n",
    "interview_builder.add_edge(\"write_section\", END)\n",
    "\n",
    "# 인터뷰 그래프 생성\n",
    "memory = MemorySaver()\n",
    "interview_graph = interview_builder.compile(checkpointer=memory).with_config(\n",
    "    run_name=\"Conduct Interviews\"\n",
    ")\n",
    "\n",
    "# 그래프 시각화\n",
    "visualize_graph(interview_graph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7bbf7293",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 분석가 목록에서 첫 번째 분석가 선택\n",
    "analysts[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76b6fe55",
   "metadata": {},
   "source": [
    "## 그래프 실행\n",
    "\n",
    "이제 그래프를 실행하고 결과를 출력합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0672f652",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Markdown\n",
    "\n",
    "# 주제 설정\n",
    "topic = \"Modular RAG 가 기존의 Naive RAG 와 어떤 차이가 있는지와 production level 에서 사용하는 이점\"\n",
    "\n",
    "# 인터뷰 시작 메시지 생성\n",
    "messages = [HumanMessage(f\"So you said you were writing an article on {topic}?\")]\n",
    "\n",
    "# 스레드 ID 설정\n",
    "config = RunnableConfig(\n",
    "    recursion_limit=100,\n",
    "    configurable={\"thread_id\": random_uuid()},\n",
    ")\n",
    "\n",
    "# 그래프 실행\n",
    "invoke_graph(\n",
    "    interview_graph,\n",
    "    {\"analyst\": analysts[0], \"messages\": messages, \"max_num_turns\": 5},\n",
    "    config,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "726ca13a",
   "metadata": {},
   "source": [
    "마크다운 형식으로 결과를 출력합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0750b48",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 완성된 인터뷰 섹션 출력\n",
    "Markdown(interview_graph.get_state(config).values[\"sections\"][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf37b34b",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(interview_graph.get_state(config).values[\"sections\"][0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cca0704d",
   "metadata": {},
   "source": [
    "## 인터뷰를 병렬로 진행 (map-reduce)\n",
    "\n",
    "- 인터뷰는 langgraph 의 `Send()` 함수를 사용하여 병렬화하며, 이는 `map` 단계에 해당합니다.\n",
    "- 인터뷰 결과는 `reduce` 단계에서 보고서 본문에 통합됩니다.\n",
    "- 최종 보고서에 서론과 결론을 작성하는 마지막 단계를 추가합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "d3dd6a95",
   "metadata": {},
   "outputs": [],
   "source": [
    "import operator\n",
    "from typing import List, Annotated\n",
    "from typing_extensions import TypedDict\n",
    "\n",
    "\n",
    "# ResearchGraphState 상태 정의\n",
    "class ResearchGraphState(TypedDict):\n",
    "    # 연구 주제\n",
    "    topic: str\n",
    "    # 생성할 분석가의 최대 수\n",
    "    max_analysts: int\n",
    "    # 인간 분석가의 피드백\n",
    "    human_analyst_feedback: str\n",
    "    # 질문을 하는 분석가 목록\n",
    "    analysts: List[Analyst]\n",
    "    # Send() API 키를 포함하는 섹션 리스트\n",
    "    sections: Annotated[list, operator.add]\n",
    "    # 최종 보고서의 서론\n",
    "    introduction: str\n",
    "    # 최종 보고서의 본문 내용\n",
    "    content: str\n",
    "    # 최종 보고서의 결론\n",
    "    conclusion: str\n",
    "    # 최종 보고서\n",
    "    final_report: str"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27a83efe",
   "metadata": {},
   "source": [
    "### LangGraph 의 Send() 함수 사용\n",
    "\n",
    "아래는 langgraph 의 `Send()` 함수를 사용하여 인터뷰를 병렬로 시작하는 함수입니다.\n",
    "\n",
    "**참고**\n",
    "\n",
    "- [LangGraph Send()](https://langchain-ai.github.io/langgraph/concepts/low_level/#send)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "6035e5ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.constants import Send\n",
    "\n",
    "\n",
    "# 모든 인터뷰를 시작\n",
    "def initiate_all_interviews(state: ResearchGraphState):\n",
    "    # 사람의 피드백 확인\n",
    "    human_analyst_feedback = state.get(\"human_analyst_feedback\")\n",
    "\n",
    "    # 만약, 사람의 피드백이 있으면 분석가 생성으로 돌아가기\n",
    "    if human_analyst_feedback:\n",
    "        return \"create_analysts\"\n",
    "\n",
    "    # 그렇지 않으면 Send() 함수를 통해 인터뷰 병렬로 시작\n",
    "    else:\n",
    "        topic = state[\"topic\"]\n",
    "        return [\n",
    "            Send(\n",
    "                \"conduct_interview\",\n",
    "                {\n",
    "                    \"analyst\": analyst,\n",
    "                    \"messages\": [\n",
    "                        HumanMessage(\n",
    "                            content=f\"So you said you were writing an article on {topic}?\"\n",
    "                        )\n",
    "                    ],\n",
    "                },\n",
    "            )\n",
    "            for analyst in state[\"analysts\"]\n",
    "        ]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d8984c7",
   "metadata": {},
   "source": [
    "### 보고서 작성 정의\n",
    "\n",
    "다음은 인터뷰 내용을 바탕으로 보고서 작성 가이드라인을 정의하고 보고서 작성 함수를 정의합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "db04abd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 보고서 작성 지시사항\n",
    "report_writer_instructions = \"\"\"You are a technical writer creating a report on this overall topic:\n",
    "\n",
    "{topic}\n",
    "\n",
    "You have a team of analysts. Each analyst has done two things:\n",
    "\n",
    "1. They conducted an interview with an expert on a specific sub-topic.\n",
    "2. They write up their finding into a memo.\n",
    "\n",
    "Your task:\n",
    "\n",
    "1. You will be given a collection of memos from your analysts.  \n",
    "2. Carefully review and analyze the insights from each memo.  \n",
    "3. Consolidate these insights into a detailed and comprehensive summary that integrates the central ideas from all the memos.  \n",
    "4. Organize the key points from each memo into the appropriate sections provided below, ensuring that each section is logical and well-structured.  \n",
    "5. Include all required sections in your report, using `### Section Name` as the header for each.  \n",
    "6. Aim for approximately 250 words per section, providing in-depth explanations, context, and supporting details.  \n",
    "\n",
    "**Sections to consider (including optional ones for greater depth):**\n",
    "\n",
    "- **Background**: Theoretical foundations, key concepts, and preliminary information necessary to understand the methodology and results.\n",
    "- **Related Work**: Overview of prior studies and how they compare or relate to the current research.\n",
    "- **Problem Definition**: A formal and precise definition of the research question or problem the paper aims to address.\n",
    "- **Methodology (or Methods)**: Detailed description of the methods, algorithms, models, data collection processes, or experimental setups used in the study.\n",
    "- **Implementation Details**: Practical details of how the methods or models were implemented, including software frameworks, computational resources, or parameter settings.\n",
    "- **Experiments**: Explanation of experimental protocols, datasets, evaluation metrics, procedures, and configurations employed to validate the methods.\n",
    "- **Results**: Presentation of experimental outcomes, often with statistical tables, graphs, figures, or qualitative analyses.\n",
    "\n",
    "To format your report:\n",
    "\n",
    "1. Use markdown formatting.\n",
    "2. Include no pre-amble for the report.\n",
    "3. Use no sub-heading.\n",
    "4. Start your report with a single title header: ## Insights\n",
    "5. Do not mention any analyst names in your report.\n",
    "6. Preserve any citations in the memos, which will be annotated in brackets, for example [1] or [2].\n",
    "7. Create a final, consolidated list of sources and add to a Sources section with the `## Sources` header.\n",
    "8. List your sources in order and do not repeat.\n",
    "\n",
    "[1] Source 1\n",
    "[2] Source 2\n",
    "\n",
    "Here are the memos from your analysts to build your report from:\n",
    "\n",
    "{context}\"\"\"\n",
    "\n",
    "\n",
    "# 보고서 작성 함수 정의\n",
    "def write_report(state: ResearchGraphState):\n",
    "    # 모든 섹션 가져오기\n",
    "    sections = state[\"sections\"]\n",
    "    topic = state[\"topic\"]\n",
    "\n",
    "    # 모든 섹션을 하나의 문자열로 연결\n",
    "    formatted_str_sections = \"\\n\\n\".join([f\"{section}\" for section in sections])\n",
    "\n",
    "    # 섹션을 요약하여 최종 보고서 작성\n",
    "    system_message = report_writer_instructions.format(\n",
    "        topic=topic, context=formatted_str_sections\n",
    "    )\n",
    "    report = llm.invoke(\n",
    "        [SystemMessage(content=system_message)]\n",
    "        + [HumanMessage(content=f\"Write a report based upon these memos.\")]\n",
    "    )\n",
    "    return {\"content\": report.content}\n",
    "\n",
    "\n",
    "# 서론과 결론 작성 지시사항\n",
    "intro_conclusion_instructions = \"\"\"You are a technical writer finishing a report on {topic}\n",
    "\n",
    "You will be given all of the sections of the report.\n",
    "\n",
    "You job is to write a crisp and compelling introduction or conclusion section.\n",
    "\n",
    "The user will instruct you whether to write the introduction or conclusion.\n",
    "\n",
    "Include no pre-amble for either section.\n",
    "\n",
    "Target around 200 words, crisply previewing (for introduction),  or recapping (for conclusion) all of the sections of the report.\n",
    "\n",
    "Use markdown formatting.\n",
    "\n",
    "For your introduction, create a compelling title and use the # header for the title.\n",
    "\n",
    "For your introduction, use ## Introduction as the section header.\n",
    "\n",
    "For your conclusion, use ## Conclusion as the section header.\n",
    "\n",
    "Here are the sections to reflect on for writing: {formatted_str_sections}\"\"\"\n",
    "\n",
    "\n",
    "# 서론 작성 함수 정의\n",
    "def write_introduction(state: ResearchGraphState):\n",
    "    # 모든 섹션 가져오기\n",
    "    sections = state[\"sections\"]\n",
    "    topic = state[\"topic\"]\n",
    "\n",
    "    # 모든 섹션을 하나의 문자열로 연결\n",
    "    formatted_str_sections = \"\\n\\n\".join([f\"{section}\" for section in sections])\n",
    "\n",
    "    # 섹션을 요약하여 서론 작성\n",
    "    instructions = intro_conclusion_instructions.format(\n",
    "        topic=topic, formatted_str_sections=formatted_str_sections\n",
    "    )\n",
    "    intro = llm.invoke(\n",
    "        [instructions] + [HumanMessage(content=f\"Write the report introduction\")]\n",
    "    )\n",
    "    return {\"introduction\": intro.content}\n",
    "\n",
    "\n",
    "# 결론 작성 함수 정의\n",
    "def write_conclusion(state: ResearchGraphState):\n",
    "    # 모든 섹션 가져오기\n",
    "    sections = state[\"sections\"]\n",
    "    topic = state[\"topic\"]\n",
    "\n",
    "    # 모든 섹션을 하나의 문자열로 연결\n",
    "    formatted_str_sections = \"\\n\\n\".join([f\"{section}\" for section in sections])\n",
    "\n",
    "    # 섹션을 요약하여 결론 작성\n",
    "    instructions = intro_conclusion_instructions.format(\n",
    "        topic=topic, formatted_str_sections=formatted_str_sections\n",
    "    )\n",
    "    conclusion = llm.invoke(\n",
    "        [instructions] + [HumanMessage(content=f\"Write the report conclusion\")]\n",
    "    )\n",
    "    return {\"conclusion\": conclusion.content}\n",
    "\n",
    "\n",
    "# 최종 보고서 작성 함수 정의\n",
    "def finalize_report(state: ResearchGraphState):\n",
    "    # 모든 섹션을 모아 최종 보고서 작성\n",
    "    content = state[\"content\"]\n",
    "    if content.startswith(\"## Insights\"):\n",
    "        content = content.strip(\"## Insights\")\n",
    "    if \"## Sources\" in content:\n",
    "        try:\n",
    "            content, sources = content.split(\"\\n## Sources\\n\")\n",
    "        except:\n",
    "            sources = None\n",
    "    else:\n",
    "        sources = None\n",
    "\n",
    "    final_report = (\n",
    "        state[\"introduction\"]\n",
    "        + \"\\n\\n---\\n\\n## Main Idea\\n\\n\"\n",
    "        + content\n",
    "        + \"\\n\\n---\\n\\n\"\n",
    "        + state[\"conclusion\"]\n",
    "    )\n",
    "    if sources is not None:\n",
    "        final_report += \"\\n\\n## Sources\\n\" + sources\n",
    "    return {\"final_report\": final_report}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c55157d",
   "metadata": {},
   "source": [
    "## 그래프 정의"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be0e8b29",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph import StateGraph\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "from langgraph.constants import START, END\n",
    "from langchain_teddynote.graphs import visualize_graph\n",
    "\n",
    "# 그래프 생성\n",
    "builder = StateGraph(ResearchGraphState)\n",
    "\n",
    "# 노드 정의\n",
    "builder.add_node(\"create_analysts\", create_analysts)\n",
    "builder.add_node(\"human_feedback\", human_feedback)\n",
    "builder.add_node(\"conduct_interview\", interview_builder.compile())\n",
    "builder.add_node(\"write_report\", write_report)\n",
    "builder.add_node(\"write_introduction\", write_introduction)\n",
    "builder.add_node(\"write_conclusion\", write_conclusion)\n",
    "builder.add_node(\"finalize_report\", finalize_report)\n",
    "\n",
    "# 엣지 정의\n",
    "builder.add_edge(START, \"create_analysts\")\n",
    "builder.add_edge(\"create_analysts\", \"human_feedback\")\n",
    "builder.add_conditional_edges(\n",
    "    \"human_feedback\", initiate_all_interviews, [\"create_analysts\", \"conduct_interview\"]\n",
    ")\n",
    "\n",
    "# 인터뷰 결과 보고서 작성\n",
    "builder.add_edge(\"conduct_interview\", \"write_report\")\n",
    "builder.add_edge(\"conduct_interview\", \"write_introduction\")\n",
    "builder.add_edge(\"conduct_interview\", \"write_conclusion\")\n",
    "\n",
    "# 보고서 최종 정리\n",
    "builder.add_edge(\n",
    "    [\"write_conclusion\", \"write_report\", \"write_introduction\"], \"finalize_report\"\n",
    ")\n",
    "builder.add_edge(\"finalize_report\", END)\n",
    "\n",
    "# 컴파일\n",
    "memory = MemorySaver()\n",
    "graph = builder.compile(interrupt_before=[\"human_feedback\"], checkpointer=memory)\n",
    "\n",
    "visualize_graph(graph)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "535c60bf",
   "metadata": {},
   "source": [
    "## 그래프 실행\n",
    "\n",
    "이제 그래프를 실행하고 결과를 확인합니다.\n",
    "\n",
    "`max_analysts`, `topic` 을 자유롭게 변경하여 실행해 보세요."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c360d65e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 입력 데이터 설정\n",
    "max_analysts = 3\n",
    "topic = \"Explain how Modular RAG differs from traditional Naive RAG and the benefits of using it at the production level.\"\n",
    "\n",
    "# config 설정\n",
    "config = RunnableConfig(\n",
    "    recursion_limit=30,\n",
    "    configurable={\"thread_id\": random_uuid()},\n",
    ")\n",
    "\n",
    "# 입력 데이터 설정\n",
    "inputs = {\"topic\": topic, \"max_analysts\": max_analysts}\n",
    "\n",
    "# 그래프 실행: 첫 번째 중단 지점까지\n",
    "invoke_graph(graph, inputs, config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed98b941",
   "metadata": {},
   "source": [
    "human_feedback 을 추가하여 분석가를 설정합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c7ee8a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 새로운 분석가 추가\n",
    "graph.update_state(\n",
    "    config,\n",
    "    {\"human_analyst_feedback\": \"Add Prof. Jeffrey Hinton as a head of AI analyst\"},\n",
    "    as_node=\"human_feedback\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "379689e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 그래프 실행\n",
    "invoke_graph(graph, None, config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c05feb4d",
   "metadata": {},
   "source": [
    "이제 사람의 피드백을 종료하고 그래프를 재개합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8d3dfb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 그래프 재개\n",
    "graph.update_state(config, {\"human_analyst_feedback\": None}, as_node=\"human_feedback\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65fd47b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 그래프 실행\n",
    "invoke_graph(graph, None, config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aca31213",
   "metadata": {},
   "source": [
    "최종 완성된 보고서를 출력합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b46b6ddc",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Markdown\n",
    "\n",
    "# 그래프의 최종 상태 가져오기\n",
    "final_state = graph.get_state(config)\n",
    "\n",
    "# 최종 보고서 가져오기\n",
    "report = final_state.values.get(\"final_report\")\n",
    "\n",
    "# 마크다운 형식으로 최종 보고서 출력\n",
    "display(Markdown(report))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4333b46d",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(report)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c12db79f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchain-kr-lwwSZlnu-py3.11",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
